import os
import shutil
from tqdm import tqdm


def cp(source_pic_root, target_pic_root):
    # 删除旧文件
    if os.path.exists(target_pic_root):
        shutil.rmtree(target_pic_root)
    os.makedirs(target_pic_root)
    # 开始复制
    for parent, _, files in os.walk(source_pic_root):
        for file in tqdm(files):
            if not file[-3:] == 'jpg':  # 不是图片则跳过
                continue
            if parent.split('/')[-1] == 'norm':
                temp_target_pic_root = os.path.join(target_pic_root, 'norm')
            else:
                temp_target_pic_root = os.path.join(target_pic_root, 'defect')
            if not os.path.exists(temp_target_pic_root):
                os.makedirs(temp_target_pic_root)
            shutil.copyfile(os.path.join(parent, file), os.path.join(temp_target_pic_root, file))


if __name__ == '__main__':
    # 1# 复制train
    cp('../data/data_for_train_balance_with_equalize_hist', '../data/data_for_bin_train_balance_with_equalize_hist')
    # 2# 复制valid
    cp('../data/data_for_valid_balance_with_equalize_hist', '../data/data_for_bin_valid_balance_with_equalize_hist')
    # 3# 复制test
    # cp('../data/data_split_test', '../data/data_split_bin_test')
